import os
import logging
from typing import List, Dict, Set, Optional
from core.models import FileInfo
from config.settings import REMOVE_DUPLICATE_DIRS, MAX_DISPLAY_DEPTH, BATCH_PROCESS_SIZE

class TreeNode:
    """树节点类"""
    def __init__(self, name: str, is_directory: bool = False):
        self.name = name
        self.is_directory = is_directory
        self.children: Dict[str, TreeNode] = {}
        self.size = 0
        self.file_count = 0

class TreeDiagramGenerator:
    """文本架构图生成器"""
    
    def __init__(self, max_depth: int = None):
        # 如果没有指定max_depth，则使用配置文件中的值
        self.max_depth = max_depth if max_depth is not None else MAX_DISPLAY_DEPTH
        self.visited_paths: Set[str] = set()  # 用于存储已访问的路径，避免重复显示
    
    def generate_tree(self, file_list: List[FileInfo], show_sizes: bool = False) -> str:
        """生成目录树状图
        
        Args:
            file_list: 文件信息列表
            show_sizes: 是否显示大小信息
        
        Returns:
            树状图字符串
        """
        # 构建目录树结构
        tree_root = self._build_tree_structure(file_list)
        
        # 生成文本树
        lines = []
        self._generate_branch(tree_root, "", True, lines, show_sizes, 0)
        
        return "\n".join(lines)
    
    def _build_tree_structure(self, file_list: List[FileInfo]) -> TreeNode:
        """构建树状结构（优化版本，支持批量处理大型文件列表）
        
        Args:
            file_list: 文件信息列表
        
        Returns:
            树的根节点
        """
        # 重置已访问路径集合
        self.visited_paths = set()
        
        # 找到根路径
        root_paths = set()
        for file_info in file_list:
            # 更准确地处理路径，支持Windows和Linux路径
            norm_path = os.path.normpath(file_info.path)
            # 获取驱动器或根目录
            if os.name == 'nt':  # Windows系统
                # Windows路径: D:\folder\file.txt -> D:
                drive = os.path.splitdrive(norm_path)[0] + '\\'
                root_paths.add(drive)
            else:  # Linux/Unix系统
                # Linux路径: /home/user/file -> /
                root_paths.add('/')
        
        # 假设第一个根路径作为根节点
        root_path = next(iter(root_paths)) if root_paths else ('D:\\' if os.name == 'nt' else '/')
        root = TreeNode(root_path, is_directory=True)
        
        # 优化：批量处理文件列表，避免一次性处理大量数据
        if file_list:
            logger = logging.getLogger('tree_generator')
            total_files = len(file_list)
            batch_size = BATCH_PROCESS_SIZE
            
            # 过滤文件列表，只保留需要处理的文件
            filtered_files = [f for f in file_list if f.is_directory or f.depth <= self.max_depth]
            
            logger.info(f"开始构建目录树，总文件数: {total_files}，实际处理文件数: {len(filtered_files)}")
            
            # 批量处理文件
            for i in range(0, len(filtered_files), batch_size):
                batch = filtered_files[i:i + batch_size]
                for file_info in batch:
                    try:
                        self._add_to_tree(root, file_info, root_path)
                    except Exception as e:
                        logger.warning(f"处理文件时出错: {file_info.path}, 错误: {e}")
                
                # 记录处理进度
                if (i + batch_size) % (batch_size * 10) == 0 or i + batch_size >= len(filtered_files):
                    processed = min(i + batch_size, len(filtered_files))
                    logger.info(f"目录树构建进度: {processed}/{len(filtered_files)} ({processed/len(filtered_files)*100:.1f}%)")
        else:
            # 如果文件列表为空，添加一个提示节点
            empty_node = TreeNode("(空目录)", is_directory=False)
            root.children["空目录"] = empty_node
        
        return root
    
    def _add_to_tree(self, root: TreeNode, file_info: FileInfo, root_path: str):
        """将文件信息添加到树中
        
        Args:
            root: 根节点
            file_info: 文件信息
            root_path: 根路径
        """
        try:
            # 获取规范化的路径
            norm_path = os.path.normpath(file_info.path)
            
            # 获取相对路径部分，使用os.path.relpath更可靠
            try:
                rel_path = os.path.relpath(norm_path, os.path.dirname(root_path))
                # 如果路径相同，relpath会返回'.'
                if rel_path == '.':
                    rel_path = ''
            except ValueError:
                # 处理不同驱动器的情况
                rel_path = norm_path[len(root_path):].strip('\\/')
            
            # 检查是否为重复路径（针对目录）
            full_path = norm_path
            if file_info.is_directory and REMOVE_DUPLICATE_DIRS:
                if full_path in self.visited_paths:
                    return  # 跳过重复目录
                self.visited_paths.add(full_path)
            
            if rel_path:
                # 同时处理Windows和Linux路径分隔符
                parts = rel_path.split(os.sep)  # 使用os.sep保证跨平台兼容性
                # 过滤掉空的路径部分
                parts = [part for part in parts if part]
                
                # 限制路径深度，只保留到max_depth层
                if len(parts) > self.max_depth:
                    parts = parts[:self.max_depth]
                
                current = root
                
                # 遍历路径部分，构建树
                for i, part in enumerate(parts):
                    is_last = i == len(parts) - 1
                    
                    if part not in current.children:
                        # 只有最后一个部分可能是文件（如果原文件不是目录）
                        is_dir = not (is_last and not file_info.is_directory and i == len(parts) - 1)
                        new_node = TreeNode(part, is_directory=is_dir)
                        # 如果是文件节点，直接设置大小
                        if not is_dir:
                            new_node.size = file_info.size
                        current.children[part] = new_node
                    
                    current = current.children[part]
                    
                    # 更新统计信息
                    if i < len(parts) - 1 or file_info.is_directory:
                        current.file_count += 1
                        if not file_info.is_directory:
                            current.size += file_info.size
            else:
                # 根目录
                root.file_count += 1
                if not file_info.is_directory:
                    root.size += file_info.size
        except Exception as e:
            # 捕获所有异常，确保即使有单个文件处理失败也不会影响整个树的构建
            import logging
            logger = logging.getLogger('tree_generator')
            logger.warning(f"处理文件路径时出错: {file_info.path}, 错误: {e}")
    
    def _generate_branch(self, node: TreeNode, prefix: str, is_last: bool,
                        lines: List[str], show_sizes: bool, current_depth: int):
        """迭代生成分支（高性能优化版本，避免递归栈溢出）
        
        Args:
            node: 当前节点
            prefix: 前缀字符串
            is_last: 是否为最后一个节点
            lines: 结果行列表
            show_sizes: 是否显示大小
            current_depth: 当前深度
        """
        # 使用栈代替递归，避免栈溢出
        stack = [(node, prefix, is_last, current_depth)]
        
        while stack:
            current_node, current_prefix, current_is_last, depth = stack.pop()
            
            # 如果超过最大深度，停止处理
            if depth > self.max_depth:
                # 在最大深度处显示省略信息
                if depth == self.max_depth + 1 and current_node.is_directory:
                    ellipsis_text = "... (以下内容省略)"
                    lines.append(current_prefix + ("    " if current_is_last else "│   ") + ellipsis_text)
                continue
            
            # 生成节点前缀
            if depth == 0:
                # 根节点没有前缀
                branch = ""
            else:
                branch = "└── " if current_is_last else "├── "
            
            # 生成节点显示文本
            node_text = current_node.name
            if current_node.is_directory:
                node_text += "/"
            
            # 添加大小和数量信息
            if show_sizes:
                if current_node.is_directory:
                    # 目录显示文件数量和大小（如果有大小）
                    if current_node.size > 0:
                        node_text += f" ({current_node.file_count} 个文件, {FileInfo.format_size(current_node.size)})"
                    else:
                        node_text += f" ({current_node.file_count} 个文件)"
                else:
                    # 文件显示大小
                    node_text += f" ({FileInfo.format_size(current_node.size)})"
            
            # 添加到结果
            lines.append(current_prefix + branch + node_text)
            
            # 处理子节点
            if current_node.is_directory and current_node.children and depth <= self.max_depth:
                # 准备子节点的前缀
                child_prefix = current_prefix
                if depth > 0:
                    child_prefix += "    " if current_is_last else "│   "
                
                # 获取子节点列表并排序（目录优先，然后按名称排序）
                children = sorted(current_node.children.items(), key=lambda x: (not x[1].is_directory, x[0]))
                
                # 限制每个目录显示的子节点数量，避免处理过多子节点导致性能问题
                max_children_to_display = 200
                total_children = len(children)
                display_children = children[:max_children_to_display]
                
                # 逆序压入栈，这样弹出时就是正序处理
                for i in reversed(range(len(display_children))):
                    name, child = display_children[i]
                    is_last_child = i == len(display_children) - 1
                    stack.append((child, child_prefix, is_last_child, depth + 1))
                
                # 如果有超过限制的子节点，显示省略信息
                if total_children > max_children_to_display:
                    remaining_count = total_children - max_children_to_display
                    if depth == 0:
                        lines.append(f"{child_prefix}... +{remaining_count} 个子项目")
                    else:
                        lines.append(f"{child_prefix}{'└── ' if current_is_last else '├── '}... +{remaining_count} 个子项目")